clear
% addpath /Applications/Dynare/4.5.6/matlab
addpath C:\Users\schorf\Dropbox\Applications\Dynare\matlab
cd('./Auxiliary Functions');

%----------------------------------------------------------------
% Set parameters
%----------------------------------------------------------------
ParaSel = 2;
Sh_size = 3;

if ParaSel == 1
   setParameters;
   mom_1_1_shift = 0;
   mom_2_1_shift = 0;
   PhiepsC_shift = 0;
   PhiepsC_addsh = 0.01;
elseif ParaSel == 2
   setParameters;
   mom_1_1_shift = -0.15;
   mom_2_1_shift = -0.15;
   PhiepsC_shift = 0;
   PhiepsC_addsh = 0.01;  
elseif ParaSel == 3
   setParameters;
   mom_1_1_shift = 0.7;
   mom_2_1_shift = 0.7;
   PhiepsC_shift = 0;
   PhiepsC_addsh = 0.01;  
end
   
%----------------------------------------------------------------
% Compute approximation tools
%----------------------------------------------------------------

% Grids
computeGrids;

% Polynomials over grids (only if using polynomials to approximate conditional expectation)
if splineOpt == 0
	computePolynomials;
end

%----------------------------------------------------------------
% Compute Steady State
%----------------------------------------------------------------

% Solve for steady state capital stock, distribution, and decision rules
coreSteadyState;

% Extract some important SS
K_star      = aggregateCapital; % in levels
mHat1_star  = mHat(1);
mHat2_star  = mHat(2);
mMean1_star = mMoments(1,1);
mMean2_star = mMoments(2,1); 

%----------------------------------------------------------------
% Save parameters in .mat files to import into Dynare 
%----------------------------------------------------------------

if splineOpt == 0	% if using polynomials to approximate individual decisions

	% Economic parameters
	save economicParameters.mat bbeta ssigma aaBar aalpha ddelta vEpsilonGrid aggEmployment ...
		mmu ttau rrhoTFP ssigmaTFP mEpsilonTransition
		
	% Approximation parameters
	save approximationParameters.mat nEpsilon nAssets nState assetsMin assetsMax nAssetsFine nStateFine nAssetsQuadrature nStateQuadrature ...
		nMeasure nMeasureCoefficients kRepSS maxIterations tolerance dampening
		
	% Grids
	save grids.mat vAssetsGridZeros vAssetsGrid mEpsilonGrid mAssetsGrid mEpsilonPrimeGrid vAssetsGridFine ...
		vAssetsGridFineZeros mEpsilonGridFine mAssetsGridFine mEpsilonPrimeGridFine vQuadratureWeights ...
		vAssetsGridQuadratureZeros vAssetsGridQuadrature mEpsilonGridQuadrature mAssetsGridQuadrature
		
	% Polynomials
	save polynomials.mat vAssetsPoly vAssetsPolySquared vAssetsPolyFine vAssetsPolyQuadrature vAssetsPolyBC
	
else	% if using splines to approximate individual decisions

	% Economic parameters
	save economicParameters.mat bbeta ssigma aaBar aalpha ddelta vEpsilonGrid aggEmployment ...
		mmu ttau rrhoTFP ssigmaTFP mEpsilonTransition
		
	% Approximation parameters
	save approximationParameters.mat nEpsilon nAssets nState assetsMin assetsMax nAssetsFine nStateFine nAssetsQuadrature nStateQuadrature ...
		nMeasure nMeasureCoefficients kRepSS maxIterations tolerance dampening
		
	% Grids
	save grids.mat vAssetsGrid mEpsilonGrid mAssetsGrid mEpsilonPrimeGrid vAssetsGridFine ...
		mEpsilonGridFine mAssetsGridFine mEpsilonPrimeGridFine vQuadratureWeights ...
		vAssetsGridQuadrature mEpsilonGridQuadrature mAssetsGridQuadrature
	
end

%----------------------------------------------------------------
% Run Dynare / Winberry Code
%----------------------------------------------------------------

if splineOpt == 0	% if using polynomials to approximate individual decisions

	dynare firstOrderDynamics_polynomials
	
else	% if using splines to approximate individual decisions

	dynare firstOrderDynamics_splines
	
end

cd('../')

save SimulDesign_data.mat mHat_1 mHat_2 moment_1_1 moment_2_1 moment_1_2 moment_2_2 moment_1_3 moment_2_3 ...
    measureCoefficient_1_1 measureCoefficient_2_1 measureCoefficient_1_2 measureCoefficient_2_2 measureCoefficient_1_3 measureCoefficient_2_3 ...
    r w aggregateTFP logAggregateOutput logAggregateInvestment logAggregateConsumption  

%%
%----------------------------------------------------------------
% Work with DYNARE output
%----------------------------------------------------------------

% extract coefficient matrices  (what's printed in the Dynare output "Policy and Transition Functions")
[ns, ne] = size(oo_.dr.ghx');
Phi1all   = zeros(ns, ne); % ns: number of state variables, ne: number of endogenous variables
Phiepsall = zeros(1, ne);
for ii = 1:ne
    Phi1all(:,ii)   = oo_.dr.ghx(oo_.dr.inv_order_var(ii),:)'; % 9x71 each column is an equation
    Phiepsall(:,ii) = oo_.dr.ghu(oo_.dr.inv_order_var(ii),:)'; % 1x71 each column is an equation
end
Phi1C   = Phi1all(:,[51 52 53 54 55 56 63 64 67])'; % each row is an equation
PhiepsC = Phiepsall(:,[51 52 53 54 55 56 63 64 67])'; % each row is an equation
% variables: mom_1_1, mom_2_1, mom_1_2, mom_2_2, mom_1_3, mom_2_3, mHat_1, mHat_2, aggTFP

Phi1mc   = Phi1all(:,[57 58 59 60 61 62])';
Phiepsmc = Phiepsall(:,[57 58 59 60 61 62])';
% variables: measureCoef_1_1, _2_1, ...

% extract steady states
ss_C  = oo_.dr.ys([51 52 53 54 55 56 63 64 67]);
ss_mc = oo_.dr.ys([57 58 59 60 61 62]);

%%
%----------------------------------------------------------------
% Adjust solution parameters to alter dynamics of the system
%----------------------------------------------------------------
Phi1C(1,3) = Phi1C(1,3) + mom_1_1_shift;
Phi1C(2,4) = Phi1C(2,4) + mom_2_1_shift;

if PhiepsC_shift == 1;
    PhiepsC_addsh_vec = PhiepsC_addsh*[1; 1; zeros(7,1)];
    PhiepsC = [PhiepsC PhiepsC_addsh_vec]
end


%%
%----------------------------------------------------------------
% IRF from Dynare output / full system
%----------------------------------------------------------------
load('./Auxiliary Functions/economicParameters.mat')

nvar_c  = 9;
nvar_r  = 2;
nvar_mc = 6;
nirf    = 60;

% Create a selection matrix to extract TFP and K
M = zeros(nvar_c,nvar_r);
M(9,1) = 1;
M(7,2) = -(1-aggEmployment)*mMean1_star;
M(1,2) = (1-aggEmployment)*(1-mHat1_star);
M(8,2) = -aggEmployment*mMean2_star;
M(2,2) = aggEmployment*(1-mHat2_star);

% TFP and Capital IRF from Model solution
IRF_c      = zeros(nvar_c,nirf);
IRF_mc     = zeros(nvar_mc,nirf);
IRF_c(:,1) = Sh_size*PhiepsC(:,1);
IRF_mc(:,1)= Sh_size*Phiepsmc(:,1);

for hh=2:nirf
    IRF_c(:,hh)  = Phi1C*IRF_c(:,hh-1);
    IRF_mc(:,hh) = Phi1mc*IRF_c(:,hh-1);
end

%----------------------------------------------------------------
% Aggregate IRF from full system
%----------------------------------------------------------------
IRF_TFP_K_c = [zeros(2,1) M'*IRF_c];
sName     = ['Para' , num2str(ParaSel) ];
dataDir   = [pwd, '/', 'Data' ,'/', sName,'/'];
[~, ~, ~] = mkdir(dataDir);
csvwrite([dataDir, 'IRF_TFP_K_c.csv'], IRF_TFP_K_c);

%----------------------------------------------------------------
% Density IRF from full system
%----------------------------------------------------------------

% we need to add steady states and we will start in the steady state
IRF_mom = [zeros(6,1) IRF_c(1:6,:)] + ss_C(1:6,1);
IRF_mc  = [zeros(6,1) IRF_mc] + ss_mc;

% asset holding grid
a_min  = 0;
a_max  = 8;
a_n    = 100;
a_grid = linspace(a_min,a_max, a_n)';

% initialize IRFs
IRF_a_employed   = zeros(a_n,nirf+1);
IRF_a_unemployed = zeros(a_n,nirf+1);
IRF_a_all        = zeros(a_n,nirf+1);

vec_percs      = [0.1; 0.2; 0.5; 0.8; 0.9];
n_percs        = length(vec_percs);
IRF_a_percs_employed   = zeros(n_percs,nirf+1);
IRF_a_percs_unemployed = zeros(n_percs,nirf+1);
IRF_a_percs_all        = zeros(n_percs,nirf+1);

% compute density for each horizon
for hh = 1:(nirf+1)
    
    sprintf('horizon= %d',hh)
    
    mom_1_1_hh = IRF_mom(1,hh);
    mom_2_1_hh = IRF_mom(2,hh);
    mom_1_2_hh = IRF_mom(3,hh);
    mom_2_2_hh = IRF_mom(4,hh);
    mom_1_3_hh = IRF_mom(5,hh);
    mom_2_3_hh = IRF_mom(6,hh);
    
    mc_1_1_hh = IRF_mc(1,hh);
    mc_2_1_hh = IRF_mc(2,hh);
    mc_1_2_hh = IRF_mc(3,hh);
    mc_2_2_hh = IRF_mc(4,hh);
    mc_1_3_hh = IRF_mc(5,hh);
    mc_2_3_hh = IRF_mc(6,hh);

    unemployed_g = @(a) exp(mc_1_1_hh * (a - mom_1_1_hh) ...
        + mc_1_2_hh * ((a - mom_1_1_hh).^2 - mom_1_2_hh) ...
        + mc_1_3_hh * ((a - mom_1_1_hh).^3 - mom_1_3_hh));
    employed_g   = @(a) exp(mc_2_1_hh * (a - mom_2_1_hh) ...
        + mc_2_2_hh * ((a - mom_2_1_hh).^2 - mom_2_2_hh) ...
        + mc_2_3_hh * ((a - mom_2_1_hh).^3 - mom_2_3_hh));
    
    unemployed_norm = integral(unemployed_g,a_min,a_max);
    employed_norm   = integral(employed_g,a_min, a_max);
    
    unemployed_pdf = @(a) unemployed_g(a)/unemployed_norm;
    employed_pdf   = @(a) employed_g(a)/employed_norm;
    all_pdf        = @(a) (1-aggEmployment)*unemployed_pdf(a) + aggEmployment*employed_pdf(a);
    
    IRF_a_employed(:,hh)   = employed_pdf(a_grid);
    IRF_a_unemployed(:,hh) = unemployed_pdf(a_grid);
    IRF_a_all(:,hh)        = all_pdf(a_grid);

    unemployed_cdf = @(a) integral(@(x) unemployed_pdf(x),a_min,a);
    employed_cdf   = @(a) integral(@(x) employed_pdf(x),a_min,a); 
    all_cdf        = @(a) integral(@(x) all_pdf(x),a_min,a); 
    
    for ii = 1:n_percs
        IRF_a_percs_employed(ii,hh)   = fzero(@(a) employed_cdf(a)-vec_percs(ii), mom_2_1_hh);
        IRF_a_percs_unemployed(ii,hh) = fzero(@(a) unemployed_cdf(a)-vec_percs(ii), mom_2_1_hh);
        IRF_a_percs_all(ii,hh)        = fzero(@(a) all_cdf(a)-vec_percs(ii), mom_2_1_hh);
    end    
        
end

sName     = ['Para' , num2str(ParaSel) ];
dataDir   = [pwd, '/', 'Data' ,'/', sName,'/'];
[~, ~, ~] = mkdir(dataDir);

csvwrite([dataDir, 'IRF_a_employed.csv'], [a_grid IRF_a_employed]);
csvwrite([dataDir, 'IRF_a_unemployed.csv'], [a_grid IRF_a_unemployed]);
csvwrite([dataDir, 'IRF_a_all.csv'], [a_grid IRF_a_all]);

csvwrite([dataDir, 'IRF_a_percs_employed.csv'], IRF_a_percs_employed);
csvwrite([dataDir, 'IRF_a_percs_unemployed.csv'], IRF_a_percs_unemployed);
csvwrite([dataDir, 'IRF_a_percs_all.csv'], IRF_a_percs_all);

% IRF for a/k density -> use change of variables
ss_K = (1-aggEmployment)*(1-mHat1_star)*mMean1_star ...
       + aggEmployment*(1-mHat2_star)*mMean2_star;
 
IRF_K                     = IRF_TFP_K_c(2,:)+ss_K;
IRF_aK_employed_vargrid   = IRF_a_employed .* IRF_K;
IRF_aK_unemployed_vargrid = IRF_a_unemployed .* IRF_K;
IRF_aK_all_vargrid        = IRF_a_all .* IRF_K;
IRF_aK_grid_vargrid       = kron( 1./IRF_K, a_grid);

IRF_aK_percs_employed   = IRF_a_percs_employed ./ IRF_K;
IRF_aK_percs_unemployed = IRF_a_percs_unemployed ./ IRF_K;
IRF_aK_percs_all        = IRF_a_percs_all ./ IRF_K;

aK_min  = 0;
aK_max  = 2.6;
aK_n    = 100;
aK_grid = linspace(0,aK_max,aK_n)';

IRF_aK_employed   = zeros(aK_n,nirf+1);
IRF_aK_unemployed = zeros(aK_n,nirf+1);
IRF_aK_all        = zeros(aK_n,nirf+1);

bw = (aK_max - aK_min)/(aK_n-1);

for hh = 1:(nirf+1)
    IRF_aK_employed(1,hh)   = IRF_aK_employed_vargrid(1,hh);
    IRF_aK_unemployed(1,hh) = IRF_aK_unemployed_vargrid(1,hh);
    IRF_aK_all(1,hh)        = IRF_aK_all_vargrid(1,hh);
    for jj = 2:aK_n;
        aK_dist      = aK_grid(jj) - IRF_aK_grid_vargrid(:,hh);
        aK_dist_sign = sign(aK_dist);
        kk = 1; 
        term_while = 0;
        while term_while == 0
            if aK_dist_sign(kk) == -1
                IRF_aK_employed(jj,hh) = IRF_aK_employed_vargrid(kk-1,hh) + ...
                  (IRF_aK_employed_vargrid(kk,hh)-IRF_aK_employed_vargrid(kk-1,hh)) / ...
                  (IRF_aK_grid_vargrid(kk,hh)-IRF_aK_grid_vargrid(kk-1,hh)) * aK_dist(kk-1);
                IRF_aK_unemployed(jj,hh) = IRF_aK_unemployed_vargrid(kk-1,hh) + ...
                  (IRF_aK_unemployed_vargrid(kk,hh)-IRF_aK_unemployed_vargrid(kk-1,hh)) / ...
                  (IRF_aK_grid_vargrid(kk,hh)-IRF_aK_grid_vargrid(kk-1,hh)) * aK_dist(kk-1);
                IRF_aK_all(jj,hh) = IRF_aK_all_vargrid(kk-1,hh) + ...
                  (IRF_aK_all_vargrid(kk,hh)-IRF_aK_all_vargrid(kk-1,hh)) / ...
                  (IRF_aK_grid_vargrid(kk,hh)-IRF_aK_grid_vargrid(kk-1,hh)) * aK_dist(kk-1);
                term_while = 1;
            elseif kk == aK_n
                term_while = 1;
            else
                kk = kk+1;
            end;
        end 
    end
end

sName     = ['Para' , num2str(ParaSel) ];
dataDir   = [pwd, '/', 'Data' ,'/', sName,'/'];
[~, ~, ~] = mkdir(dataDir);

csvwrite([dataDir, 'IRF_aK_employed.csv'], [aK_grid IRF_aK_employed]);
csvwrite([dataDir, 'IRF_aK_unemployed.csv'], [aK_grid IRF_aK_unemployed]);
csvwrite([dataDir, 'IRF_aK_all.csv'], [aK_grid IRF_aK_all]);

csvwrite([dataDir, 'IRF_aK_percs_employed.csv'], IRF_aK_percs_employed);
csvwrite([dataDir, 'IRF_aK_percs_unemployed.csv'], IRF_aK_percs_unemployed);
csvwrite([dataDir, 'IRF_aK_percs_all.csv'], IRF_aK_percs_all);

%%
%----------------------------------------------------------------
% Plot density (differential) response
%----------------------------------------------------------------

sName   = ['Para' , num2str(ParaSel) ];
figDir  = [pwd, '/', 'Figures' ,'/', sName,'/'];
[~, ~, ~] = mkdir(figDir);

% Figure: 
hh1 = 2;
hh2 = 11;
figure(1);clf;
set(figure(1),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(a_grid,(IRF_a_employed(:,hh1)-IRF_a_employed(:,1)),'Color','b','LineStyle','-','LineWidth',2)  
hold on
plot(a_grid,(IRF_a_employed(:,hh2)-IRF_a_employed(:,1)),'Color','r','LineStyle',':','LineWidth',2)  
hold on
plot(a_grid, zeros(a_n), 'Color', 'k', 'LineStyle', '-', 'LineWidth',2)
set(gca,'FontSize',50)
xlim([0 a_max])

sFigName = [sName, '_IRF_AggSh1_a_DensityDiff_Employed.pdf'];
saveas(figure(1), [figDir sFigName] );

% Figure: 
hh1 = 4;
figure(2);clf;
set(figure(2),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(aK_grid,IRF_aK_employed(:,hh1),'Color','b','LineStyle','-','LineWidth',2)  
hold on
plot(IRF_aK_grid_vargrid(:,hh1), IRF_aK_employed_vargrid(:,hh1),'Color','r','LineStyle',':','LineWidth',2)  
hold on
plot(aK_grid, zeros(aK_n), 'Color', 'k', 'LineStyle', '-', 'LineWidth',2)
set(gca,'FontSize',50)
xlim([0 aK_max])

sFigName = [sName, '_IRF_AggSh1_aK_Density_Employed_Interpolate.pdf'];
saveas(figure(2), [figDir sFigName] );

% Figure: 
hh1 = 2;
hh2 = 11;
figure(3);clf;
set(figure(3),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(aK_grid,(IRF_aK_employed(:,hh1)-IRF_aK_employed(:,1)),'Color','b','LineStyle','-','LineWidth',2)  
hold on
plot(aK_grid,(IRF_aK_employed(:,hh2)-IRF_aK_employed(:,1)),'Color','r','LineStyle',':','LineWidth',2)  
hold on
plot(aK_grid, zeros(aK_n), 'Color', 'k', 'LineStyle', '-', 'LineWidth',2)
set(gca,'FontSize',50)
xlim([0 aK_max])

sFigName = [sName, '_IRF_AggSh1_aK_DensityDiff_Employed.pdf'];
saveas(figure(3), [figDir sFigName] );

%%
%----------------------------------------------------------------
% Plot percentile response
%----------------------------------------------------------------

% Figure: Percentile IRF to a TFP shock
figure(3);clf;
set(figure(3),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(1:(nirf),IRF_aK_percs_all(1,2:end)'-IRF_aK_percs_all(1,1),'Color','r','LineStyle',':','LineWidth',2)  
hold on
plot(1:(nirf),IRF_aK_percs_all(2,2:end)'-IRF_aK_percs_all(2,1),'Color','g','LineStyle','--','LineWidth',2)  
hold on
plot(1:(nirf),IRF_aK_percs_all(3,2:end)'-IRF_aK_percs_all(3,1),'Color','b','LineStyle','-','LineWidth',2)  
hold on
plot(1:(nirf),IRF_aK_percs_all(4,2:end)'-IRF_aK_percs_all(4,1),'Color','g','LineStyle','--','LineWidth',2)  
hold on
plot(1:(nirf),IRF_aK_percs_all(5,2:end)'-IRF_aK_percs_all(5,1),'Color','r','LineStyle',':','LineWidth',2)  
hold on
set(gca,'FontSize',50)
xlim([0 nirf])

sFigName = [sName, '_IRF_ak_percs_all.pdf'];
saveas(figure(3), [figDir sFigName] );


%%
%----------------------------------------------------------------
% IRF from population regression / reduced system
%----------------------------------------------------------------

% Compute population acf of S_r
S_c_Gam0 = inv(eye(length(Phi1C)^2) - kron(Phi1C,Phi1C))*vec(PhiepsC*PhiepsC');
S_c_Gam0 = reshape(S_c_Gam0,length(Phi1C),length(Phi1C));
S_c_Gam1 = Phi1C*S_c_Gam0;

% Derive population acf of S_r
S_r_Gam0 = M'*S_c_Gam0*M;
S_r_Gam1 = M'*S_c_Gam1*M;

% Population VAR1 regression
S_r_Phi1 = (S_r_Gam0 \ S_r_Gam1')';

% Population cov matrix for S_r innovation derived from S_r VAR
S_r_Sigma = S_r_Gam0 - S_r_Gam1*(S_r_Gam0 \ S_r_Gam1');

%"Innovation covariance matrix for (TFP, K), full system"
%M'*PhiepsC*PhiepsC'*M
%"Innovation covariance matrix for (TFP, K), reduced system"
%S_r_Sigma

% Impact matrix for IRFs
S_r_Sigma_chol = chol(S_r_Sigma,'lower');

% TFP and Capital IRF from Model solution
IRF_TFP_K_r = zeros(nvar_r,nirf);
IRF_TFP_K_r(:,1) = Sh_size*S_r_Sigma_chol(:,1);
for hh=2:nirf
    IRF_TFP_K_r(:,hh) = S_r_Phi1*IRF_TFP_K_r(:,hh-1);
end
IRF_TFP_K_r = [ zeros(2,1) IRF_TFP_K_r];

sName     = ['Para' , num2str(ParaSel) ];
dataDir   = [pwd, '/', 'Data' ,'/', sName,'/'];
[~, ~, ~] = mkdir(dataDir);
csvwrite([dataDir, 'IRF_TFP_K_r.csv'], IRF_TFP_K_r);


%%
%----------------------------------------------------------------
% Plot IRFs
%----------------------------------------------------------------

sName  = ['Para' , num2str(ParaSel) ];
figDir = [pwd, '/', 'Figures' ,'/', sName,'/'];
[~, ~, ~] = mkdir(figDir);

% Figure: TFP IRF to a TFP shock
figure(4);clf;
set(figure(4),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(1:nirf,IRF_TFP_K_c(1,2:nirf+1)','Color','b','LineStyle','-','LineWidth',4)  
hold on
plot(1:nirf,IRF_TFP_K_r(1,2:nirf+1)','Color','r','LineStyle',':','LineWidth',4)  
hold on
plot(1:nirf, zeros(nirf), 'Color', 'k', 'LineStyle', '-', 'LineWidth',2)
set(gca,'FontSize',50)
xlim([1 nirf])

sFigName = [sName, '_IRF_AggSh1_TFP.pdf'];
saveas(figure(4), [figDir sFigName] );
% close all

% Figure: K IRF to a TFP shock
figure(5);clf;
set(figure(5),'PaperType','usletter','PaperOrientation','Landscape','PaperPosition',[0.1 0.1 11 8.5]);
plot(1:nirf,IRF_TFP_K_c(2,2:nirf+1)','Color','b','LineStyle','-','LineWidth',4)  
hold on
plot(1:nirf,IRF_TFP_K_r(2,2:nirf+1)','Color','r','LineStyle',':','LineWidth',4)  
hold on
plot(1:nirf, zeros(nirf), 'Color', 'k', 'LineStyle', '-', 'LineWidth',2)
set(gca,'FontSize',50)
xlim([1 nirf])

sFigName = [sName, '_IRF_AggSh1_K.pdf'];
saveas(figure(5), [figDir sFigName] );
% close all

